import numpy as np
from lstm import LstmParam, LstmNetwork
import matplotlib.pyplot as plt


class ToyLossLayer:
    """
    Computes square loss with first element of hidden layer array.
    """

    @classmethod
    def loss(self, pred, label):
        return (pred[0] - label) ** 2

    @classmethod
    def bottom_diff(self, pred, label):
        diff = np.zeros_like(pred)
        diff[0] = 2 * (pred[0] - label)
        return diff


def example_0():
    # learns to repeat simple sequence from random inputs
    np.random.seed(0)

    # parameters for input data dimension and lstm cell count
    mem_cell_ct = 100
    x_dim = 50
    lstm_param = LstmParam(mem_cell_ct, x_dim)
    lstm_net = LstmNetwork(lstm_param)
    y_list = [-0.5, 0.2, 0.1, -0.5]

    y_list = [-0.43, -0.46, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.41, -0.35, -0.27, -0.18, -0.07, 0.12,
              0.35, 0.17,
              0.04, 0.06, 0.01, 0.36, 0.5, -0.38]

    input_val_arr = [np.random.random(x_dim) for _ in y_list]
    print(len(input_val_arr))
    print(input_val_arr)

    loss_list = []
    pre_list = []

    for cur_iter in range(1000):
        print("iter", "%2s" % str(cur_iter), end=": ")
        for ind in range(len(y_list)):
            lstm_net.x_list_add(input_val_arr[ind])

        print("y_pred = [" +
              ", ".join(["% 2.5f" % lstm_net.lstm_node_list[ind].state.h[0] for ind in range(len(y_list))]) +
              "]", end=", ")

        pre_list = [lstm_net.lstm_node_list[ind].state.h[0] for ind in range(len(y_list))]

        loss = lstm_net.y_list_is(y_list, ToyLossLayer)
        loss_list.append(loss)
        print("loss:", "%.3e" % loss)
        lstm_param.apply_diff(lr=0.1)
        lstm_net.x_list_clear()

    x = range(1, len(loss_list) + 1)
    plt.plot(x, loss_list)
    plt.show()

    print("=======================")

    max_value = 81

    print(pre_list)
    print("=======================")
    for i in range(len(pre_list)):
        pre_list[i] = round(pre_list[i], 2)
        pre_list[i] = (pre_list[i] + 0.5) * max_value
        pre_list[i] = round(pre_list[i], 1)
    print("预测：", pre_list)

    for i in range(len(y_list)):
        y_list[i] = (y_list[i] + 0.5) * max_value
        y_list[i] = round(y_list[i], 1)
    print("实际：", y_list)

    delta = []
    for i in range(len(y_list)):
        num = pre_list[i] - y_list[i]
        delta.append(round(num, 1))
    print("相差：", delta)

    x = range(1, len(y_list) + 1)
    plt.plot(x, pre_list, color='red', label='pre')
    plt.plot(x, y_list, color='green', label='actual')
    plt.legend()  # 显示图例
    plt.xlabel('iteration times')
    plt.ylabel('value')

    plt.show()


if __name__ == "__main__":
    example_0()
